#pragma once

#include "OctreePublic.h"
#include <vector>

/** A concise iteration over the children of an octree node. */
#define FOREACH_OCTREE_CHILD_NODE(ChildRef) \
	for(FOctreeChildNodeRef ChildRef(0);!ChildRef.IsNULL();ChildRef.Advance())

/** A reference to a child of an octree node. */
class FOctreeChildNodeRef
{
public:

	union
	{
		struct
		{
			unsigned long X : 1;
			unsigned long Y : 1;
			unsigned long Z : 1;
			unsigned long bNULL : 1;
		};
		unsigned long Index : 3;
	};

	/** Initialization constructor. */
	FOctreeChildNodeRef(int InX,int InY,int InZ)
	:	X(InX)
	,	Y(InY)
	,	Z(InZ)
	,	bNULL(false)
	{}

	/** Initialized the reference with a child index. */
	FOctreeChildNodeRef(int InIndex = 0)
	:	bNULL(false)
	,	Index(InIndex)
	{}

	/** Advances the reference to the next child node.  If this was the last node remain, sets bInvalid=true. */
	__forceinline void Advance()
	{
		if(Index < 7)
		{
			++Index;
		}
		else
		{
			bNULL = true;
		}
	}

	/** @return true if the reference isn't set. */
	__forceinline bool IsNULL() const
	{
		return bNULL;
	}
};

/** A subset of an octree node's children that intersect a bounding box. */
class FOctreeChildNodeSubset
{
public:

	union
	{
		struct 
		{
			unsigned long bPositiveX : 1;
			unsigned long bPositiveY : 1;
			unsigned long bPositiveZ : 1;
			unsigned long bNegativeX : 1;
			unsigned long bNegativeY : 1;
			unsigned long bNegativeZ : 1;
		};

		struct
		{
			/** Only the bits for the children on the positive side of the splits. */
			unsigned long PositiveChildBits : 3;

			/** Only the bits for the children on the negative side of the splits. */
			unsigned long NegativeChildBits : 3;
		};

		/** All the bits corresponding to the child bits. */
		unsigned long ChildBits : 6;

		/** All the bits used to store the subset. */
		unsigned long AllBits;
	};

	/** Initializes the subset to be empty. */
	FOctreeChildNodeSubset()
	:	AllBits(0)
	{}

	/** Initializes the subset to contain a single node. */
	FOctreeChildNodeSubset(FOctreeChildNodeRef ChildRef)
	:	AllBits(0)
	{
		// The positive child bits correspond to the child index, and the negative to the NOT of the child index.
		PositiveChildBits = ChildRef.Index;
		NegativeChildBits = ~ChildRef.Index;
	}

	/** Determines whether the subset contains a specific node. */
	bool Contains(FOctreeChildNodeRef ChildRef) const;
};

/** The context of an octree node, derived from the traversal of the tree. */
class FOctreeNodeContext
{
public:

	/** The node bounds are expanded by their extent divided by LoosenessDenominator. */
	enum { LoosenessDenominator = 16 };

	/** The bounds of the node. */
	BoxCenterAndExtent Bounds;

	/** The extent of the node's children. */
	float ChildExtent;

	/** The offset of the childrens' centers from the center of this node. */
	float ChildCenterOffset;

	/** Default constructor. */
	FOctreeNodeContext()
	{}

	/** Initialization constructor. */
	FOctreeNodeContext(const BoxCenterAndExtent& InBounds)
	:	Bounds(InBounds)
	{
		// A child node's tight extents are half its parent's extents, and its loose extents are expanded by 1/LoosenessDenominator.
		const float TightChildExtent = Bounds.Extent.X * 0.5f;
		const float LooseChildExtent = TightChildExtent * (1.0f + 1.0f / (float)LoosenessDenominator);

		ChildExtent = LooseChildExtent;
		ChildCenterOffset = Bounds.Extent.X - LooseChildExtent;
	}

	/** Child node initialization constructor. */
	__forceinline FOctreeNodeContext GetChildContext(FOctreeChildNodeRef ChildRef) const
	{
		return FOctreeNodeContext(BoxCenterAndExtent(
				Vector3(
					Bounds.Center.X + ChildCenterOffset * (-1.0f + 2 * ChildRef.X),
					Bounds.Center.Y + ChildCenterOffset * (-1.0f + 2 * ChildRef.Y),
					Bounds.Center.Z + ChildCenterOffset * (-1.0f + 2 * ChildRef.Z)),
				Vector3(
					ChildExtent,
					ChildExtent,
					ChildExtent)
				));
	}
	
	/**
	 * Determines which of the octree node's children intersect with a bounding box.
	 * @param BoundingBox - The bounding box to check for intersection with.
	 * @return A subset of the children's nodes that intersect the bounding box.
	 */
	FOctreeChildNodeSubset GetIntersectingChildren(const BoxCenterAndExtent& BoundingBox) const;

	/**
	 * Determines which of the octree node's children contain the whole bounding box, if any.
	 * @param BoundingBox - The bounding box to check for intersection with.
	 * @return The octree's node that the bounding box is farthest from the outside edge of, or an invalid node ref if it's not contained
	 *			by any of the children.
	 */
	FOctreeChildNodeRef GetContainingChild(const BoxCenterAndExtent& BoundingBox) const;
};

/** An octree. */
template<typename ElementType,typename OctreeSemantics>
class TOctree
{
public:

	typedef std::vector<ElementType> ElementArrayType;
	typedef typename ElementArrayType::iterator ElementIt;
	typedef typename ElementArrayType::const_iterator ElementConstIt;

	/** A node in the octree. */
	class FNode
	{
	public:

		friend class TOctree;

		/** Initialization constructor. */
		FNode(const FNode* InParent)
		:	Parent(InParent)
		,	InclusiveNumElements(0)
		,	bIsLeaf(true)
		{
			FOREACH_OCTREE_CHILD_NODE(ChildRef)
			{
				Children[ChildRef.Index] = NULL;
			}
		}

		/** Destructor. */
		~FNode()
		{
			FOREACH_OCTREE_CHILD_NODE(ChildRef)
			{
				delete Children[ChildRef.Index];
			}
		}

		// Accessors.
		__forceinline const ElementArrayType& GetElements() const { return Elements; }
		__forceinline bool IsLeaf() const { return bIsLeaf; }
		__forceinline bool HasChild(FOctreeChildNodeRef ChildRef) const
		{
			return Children[ChildRef.Index] != NULL && Children[ChildRef.Index]->InclusiveNumElements > 0;
		}
		__forceinline FNode* GetChild(FOctreeChildNodeRef ChildRef) const
		{
			return Children[ChildRef.Index];
		}

	private:
		/** The elements in this node. */
		mutable ElementArrayType Elements;

		/** The parent of this node. */
		const FNode* Parent;

		/** The children of the node. */
		mutable FNode* Children[8];

		/** The number of elements contained by the node and its child nodes. */
		mutable unsigned long InclusiveNumElements : 31;

		/** true if the meshes should be added directly to the node, rather than subdividing when possible. */
		mutable unsigned long bIsLeaf : 1;
	};



	/** A reference to an octree node, its context, and a read lock. */
	class FNodeReference
	{
	public:

		const FNode* Node;
		FOctreeNodeContext Context;

		/** Default constructor. */
		FNodeReference():
			Node(NULL),
			Context()
		{}

		/** Initialization constructor. */
		FNodeReference(const FNode* InNode,const FOctreeNodeContext& InContext):
			Node(InNode),
			Context(InContext)
		{}
	};	

	/** The default iterator allocator gives the stack enough inline space to contain a path and its siblings from root to leaf. */
	// typedef TInlineAllocator<7 * (14 - 1) + 8> DefaultStackAllocator;

	/** An octree node iterator. */
	template <int MaxNum = 7 * (14 - 1) + 8>
	class NodeIterator
	{
	public:
		NodeIterator(const TOctree& Tree) : NodeStack(0), CurrentNode(FNodeReference(&Tree.RootNode,Tree.RootNodeContext)) {}
		NodeIterator(const FNode& Node,const FOctreeNodeContext& Context) : NodeStack(0), CurrentNode(FNodeReference(&Node,Context)) {}

		void PushChild(FOctreeChildNodeRef ChildRef)
		{
			if (NodeStack >= MaxNum)
				return;

			NodePool[NodeStack++] = FNodeReference(
				CurrentNode.Node->GetChild(ChildRef),
				CurrentNode.Context.GetChildContext(ChildRef));
		}

		/** Iterates to the next node. */
		void Advance()
		{
			if (NodeStack > 0)
				CurrentNode = NodePool[--NodeStack];
			else
				CurrentNode = FNodeReference();
		}

		bool HasPendingNodes() const { return CurrentNode.Node != NULL; }
		const FNode& GetCurrentNode() const { return *CurrentNode.Node; }
		const FOctreeNodeContext& GetCurrentContext() const { return CurrentNode.Context; }

	private:
		int NodeStack;
		FNodeReference CurrentNode;
		FNodeReference NodePool[MaxNum];
	};
	typedef NodeIterator<> TConstIterator;

	/** Iterates over the elements in the octree that intersect a bounding box. */
	class TConstElementBoxIterator
	{
	public:

		/** Iterates to the next element. */
		void Advance()
		{
			++ElementIt;
			AdvanceToNextIntersectingElement();
		}

		/** Checks if there are any elements left to iterate over. */
		bool HasPendingElements() const
		{
			return NodeIt.HasPendingNodes();
		}

		/** Initialization constructor. */
		TConstElementBoxIterator(const TOctree& Tree,const BoxCenterAndExtent& InBoundingBox)
		:	IteratorBounds(InBoundingBox)
		,	NodeIt(Tree)
		,	ElementIt(Tree.RootNode.GetElementIt())
		{
			ProcessChildren();
			AdvanceToNextIntersectingElement();
		}

		// Accessors.
		const ElementType& GetCurrentElement() const
		{
			return *ElementIt;
		}

	private:

		/** The bounding box to check for intersection with. */
		BoxCenterAndExtent IteratorBounds;

		/** The octree node iterator. */
		TConstIterator NodeIt;

		/** The element iterator for the current node. */
		ElementConstIt ElementIt;

		/** Processes the children of the current node. */
		void ProcessChildren()
		{
			// Add the child nodes that intersect the bounding box to the node iterator's stack.
			const FNode& CurrentNode = NodeIt.GetCurrentNode();
			const FOctreeNodeContext& Context = NodeIt.GetCurrentContext();
			const FOctreeChildNodeSubset IntersectingChildSubset = Context.GetIntersectingChildren(IteratorBounds);
			FOREACH_OCTREE_CHILD_NODE(ChildRef)
			{
				if(IntersectingChildSubset.Contains(ChildRef) && CurrentNode.HasChild(ChildRef))
				{
					NodeIt.PushChild(ChildRef);
				}
			}
		}

		/** Advances the iterator to the next intersecting primitive, starting at a primitive in the current node. */
		void AdvanceToNextIntersectingElement()
		{
			// Keep trying elements until we find one that intersects or run out of elements to try.
			while(NodeIt.HasPendingNodes())
			{
				// Check if we've advanced past the elements in the current node.
				if(ElementIt)
				{
					// Check if the current element intersects the bounding box.
					if(Intersect(OctreeSemantics::GetBoundingBox(*ElementIt),IteratorBounds))
					{
						// If it intersects, break out of the advancement loop.
						break;
					}
					else
					{
						// If it doesn't intersect, skip to the next element.
						++ElementIt;
					}
				}
				else
				{
					// Advance to the next node.
					NodeIt.Advance();
					if(NodeIt.HasPendingNodes())
					{
						ProcessChildren();

						// The element iterator can't be assigned to, but it can be replaced by Move.
						Move(ElementIt,NodeIt.GetCurrentNode().GetElementIt());
					}
				}
			};
		}
	};

	/**
	 * Adds an element to the octree.
	 * @param Element - The element to add.
	 * @return An identifier for the element in the octree.
	 */
	void AddElement(ElementType& Element);

	/**
	 * Removes an element from the octree.
	 * @param ElementId - The element to remove from the octree.
	 */
	void RemoveElement(FOctreeElementId ElementId);

	/** Accesses an octree element by ID. */
	ElementType& GetElementById(FOctreeElementId ElementId);

	/** Writes stats for the octree to the log. */
	void DumpStats();

	/** Initialization constructor. */
	TOctree(const Vector3& InOrigin, float InExtent);

private:

	/** The octree's root node. */
	FNode RootNode;

	/** The octree's root node's context. */
	FOctreeNodeContext RootNodeContext;

	/** The extent of a leaf at the maximum allowed depth of the tree. */
	float MinLeafExtent;

	/** Adds an element to a node or its children. */
	void AddElementToNode(
		ElementType& Element,
		const FNode& InNode,
		const FOctreeNodeContext& InContext);
};

#include "Octree.inl"
